# !/usr/bin/env python
# -*- coding: utf-8 -*-
# @File  : 手写数字识别_KNN实现分类.py
# @Author: dongguangwen
# @Date  : 2025-01-19 20:05
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
import joblib


def train_model():
    # 1.加载数据
    data = pd.read_csv('./data/手写数字识别.csv')

    # 2.数据预处理-归一化
    x = data.iloc[:, 1:] / 255
    y = data.iloc[:, 0]

    # 3.分割数据集
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, stratify=y, random_state=22)

    # 4.模型训练
    estimator = KNeighborsClassifier(n_neighbors=3)
    estimator.fit(x_train, y_train)

    # 5.模型评估
    acc = estimator.score(x_test, y_test)
    print(f'Accuracy is {acc}')

    # 6.模型保存
    joblib.dump(estimator, './data/knn_model.pth')


def model_load():
    # 7.模型加载
    model = joblib.load('./data/knn_model.pth')
    return model


def model_predict(model):
    # 8.模型预测
    img = plt.imread('./data/digit.png')
    # plt.imshow(img, cmap='gray')
    img = img[:, :, 1].reshape(1, -1)
    pred = model.predict(img)
    print(f'Predicted: {pred}')  # Predicted: [0]


if __name__ == '__main__':
    train_model()
    model = model_load()
    model_predict(model)
